from PySide6 import QtGui, QtWidgets, QtCore
from common.globaldefines import GetTimeStamp, g_DebugPrintHook
import common.globaldefines
import application
import scrcpy
import labelmanager
import os
import time

def convertOriginal2Display(original_rect, original_size, display_size):
    """
    将原图坐标转换到显示区域坐标
    
    Args:
        original_rect: 原图中的矩形 (x, y, width, height)
        original_size: 原图尺寸 (width, height)
        display_size: 显示区域尺寸 (width, height)
    
    Returns:
        display_rect: 显示区域中的矩形 (x, y, width, height)
    """
    original_x, original_y, original_width, original_height = original_rect
    original_width_total, original_height_total = original_size
    display_width, display_height = display_size
    
    # 计算缩放比例
    scale_x = display_width / original_width_total  # 800 / 1280
    scale_y = display_height / original_height_total  # 480 / 720
    
    # 转换坐标和尺寸
    display_x = int(original_x * scale_x)
    display_y = int(original_y * scale_y)
    display_width = int(original_width * scale_x)
    display_height = int(original_height * scale_y)
    
    return (display_x, display_y, display_width, display_height)

def convertDisplay2Original(display_rect, display_size, original_size):
    """
    将显示区域的坐标转换回原图坐标
    
    Args:
        display_rect: 显示区域中的矩形 (x, y, width, height)
        display_size: 显示区域尺寸 (width, height)
        original_size: 原图尺寸 (width, height)
    
    Returns:
        original_rect: 原图中的矩形 (x, y, width, height)
    """
    display_x, display_y, display_width, display_height = display_rect
    display_width_total, display_height_total = display_size
    original_width, original_height = original_size
    
    # 计算缩放比例
    scale_x = original_width / display_width_total
    scale_y = original_height / display_height_total
    
    # 转换坐标和尺寸
    original_x = int(display_x * scale_x)
    original_y = int(display_y * scale_y)
    original_width = int(display_width * scale_x)
    original_height = int(display_height * scale_y)
    
    return (original_x, original_y, original_width, original_height)

def add_black_mask(image: QtGui.QImage, opacity: float = 0.5) -> QtGui.QImage:
    """
    为QImage添加黑色半透明蒙版
    :param image: 原始QImage对象
    :param opacity: 透明度(0-1)，1为完全不透明
    :return: 带蒙版的新QImage对象
    """
    if image.isNull():
        return image
        
    # 创建副本避免修改原图
    masked_image = image.copy()
    
    # 创建QPainter进行绘制
    painter = QtGui.QPainter(masked_image)
    try:
        # 设置半透明黑色画刷
        painter.setBrush(QtGui.QColor(0, 0, 0, int(255 * opacity)))
        painter.setPen(QtCore.Qt.NoPen)
        
        # 绘制覆盖整个图像的矩形
        painter.drawRect(masked_image.rect())
    finally:
        painter.end()
        
    return masked_image

class DisplayWidget(QtWidgets.QWidget):

    signalNewLabel = QtCore.Signal(str)

    def __init__(self):
        super().__init__()
        self.init()
        self._image: QtGui.QImage = None
        self._darkimage: QtGui.QPixmap = None
        self._size = QtCore.QSize()
        self._dsize = QtCore.QSize()
        self._beginMousePos = None
        self._endMousePos = None
        self.setReadOnlyMode(True)
        self._lastSelectedArea: QtCore.Qt.QRect = None
        self._historySelectArea = []
        self._title = ""

    def init(self):
        pass

    def updateScreen(self, img: QtGui.QImage):
        self._image = img
        self._size = img.size()
        self.update()

    def resizeEvent(self, event):
        return super().resizeEvent(event)
    
    def setReadOnlyMode(self, b):
        self._readonlymode = b

    def paintEvent(self, event):
        painter = QtGui.QPainter(self)

        painter.setBrush(QtGui.QColor(120, 150, 200))
        painter.setPen(QtGui.QColor(0, 0, 0))
        painter.drawRect(self.rect())
        if self._image:
            if self._beginMousePos:
                if not self._darkimage:
                    self._darkimage = add_black_mask(self._image)
                bgimg = self._darkimage
            else:
                bgimg = self._image
            img = bgimg.scaled(self.width(), self.height()
                                     , QtCore.Qt.AspectRatioMode.IgnoreAspectRatio
                                     , QtCore.Qt.TransformationMode.SmoothTransformation)
            self._dsize = img.size()
            painter.drawImage(QtCore.QPoint(0, 0), img)
            painter.drawImage(0, 0, img)
            r = self.drawSelectArea(painter)
            painter.setBrush(QtCore.Qt.NoBrush)
            if not r:
                # 设置绿色虚线画笔
                dash_pen = QtGui.QPen(QtGui.QColor(44, 255, 44))  # 绿色
                dash_pen.setStyle(QtCore.Qt.DashLine)
                dash_pen.setWidth(2)
        
                # 绘制虚线矩形（边框）
                painter.setPen(dash_pen)
                for info in labelmanager.g_Instance._labels.get(self._title, []):
                    x, y, w, h = convertOriginal2Display(info[1:]
                                            , (self._image.width(), self._image.height())
                                            , (self.width(), self.height()))
                    nrect = QtCore.QRect(x, y, w, h)
                    painter.drawRect(nrect)
                    painter.drawText(x, y, info[0])
            #self.setFixedSize(self._dsize)
        return super().paintEvent(event)

    def transQMousePosToDevicePos(self, p):
        cx, cy = p.x(), p.y()
        if cx < 0: cx = 0
        if cx > self._dsize.width(): cx = self._dsize.width()
        if cy < 0: cy = 0
        if cy > self._dsize.height(): cy = self._dsize.height()
        wr = self._size.width() / self._dsize.width()
        hr = self._size.height() / self._dsize.height()
        realx, realy = cx * wr, cy * hr
        return realx, realy

    def mousePressEvent(self, event):
        realx, realy = self.transQMousePosToDevicePos(event.pos())
        #DebugPrint("press:", realx, realy)
        if not self._readonlymode:
            self._beginMousePos = event.pos()
        else:
            application.Application.instance().pushMessage("scrcpy", "mousepress", realx, realy, scrcpy.ACTION_DOWN)
        return super().mousePressEvent(event)
    
    def mouseReleaseEvent(self, event):
        realx, realy = self.transQMousePosToDevicePos(event.pos())
        self._beginMousePos = None
        self._endMousePos = None
        #DebugPrint("release:", realx, realy)
        if self._readonlymode:
            application.Application.instance().pushMessage("scrcpy", "mouserelease", realx, realy, scrcpy.ACTION_UP)
        else:
            # 鼠标松开之后，把框选加入历史列表里
            self._historySelectArea.append(self._lastSelectedArea)
            inputretry = True
            title = "Set label name"
            txt = "Label:"
            while inputretry:
                name = QtWidgets.QInputDialog().getText(self, title, txt)
                if name[1]:
                    DebugPrint(name[0])
                    if not labelmanager.g_Instance.checkValidName(self._title, name[0]):
                        txt = f"\"{name[0]}\" 已存在，换一个吧"
                        continue
                    labelmanager.g_Instance.addLabel(self._title, name[0]
                                                    , self._lastSelectedArea.x()
                                                    , self._lastSelectedArea.y()
                                                    , self._lastSelectedArea.width()
                                                    , self._lastSelectedArea.height())
                    self.signalNewLabel.emit(name[0])
                    inputretry = False
                    break
                else:
                    inputretry = False
        return super().mouseReleaseEvent(event)
    
    def mouseMoveEvent(self, event):
        realx, realy = self.transQMousePosToDevicePos(event.pos())
        #DebugPrint("move:", realx, realy)
        if not self._readonlymode:
            self._endMousePos = event.pos()
        else:
            application.Application.instance().pushMessage("scrcpy", "mousemove", realx, realy, scrcpy.ACTION_MOVE)
        return super().mouseMoveEvent(event)

    def getSelectArea(self):
        if self._beginMousePos and self._endMousePos:
            if self._beginMousePos.x() > self._endMousePos.x() and self._beginMousePos.y() > self._endMousePos.y():
                startpos = (self._endMousePos.x(), self._endMousePos.y())
                endpos = (self._beginMousePos.x(), self._beginMousePos.y())
            elif self._beginMousePos.x() < self._endMousePos.x() and self._beginMousePos.y() > self._endMousePos.y():
                startpos = (self._beginMousePos.x(), self._endMousePos.y())
                endpos = (self._endMousePos.x(), self._beginMousePos.y())
            elif self._beginMousePos.x() > self._endMousePos.x() and self._beginMousePos.y() < self._endMousePos.y():
                startpos = (self._endMousePos.x(), self._beginMousePos.y())
                endpos = (self._beginMousePos.x(), self._endMousePos.y())
            else:
                startpos = (self._beginMousePos.x(), self._beginMousePos.y())
                endpos = (self._endMousePos.x(), self._endMousePos.y())
            # 要转成图片的大小记录下来
            w = endpos[0] - startpos[0] 
            h = endpos[1] - startpos[1]
            wscale = self.size().width() / self._image.size().width()
            hscale = self.size().height() / self._image.size().height()
            wscale = 1
            hscale = 1
            selectrect: QtCore.QRect = QtCore.QRect(int(startpos[0] / wscale)
                                       , int(startpos[1] / hscale)
                                       , int(w / wscale), int(h / hscale))
            return selectrect
        return None

    def drawSelectArea(self, p: QtGui.QPainter):
        if not self._beginMousePos:
            return
        if not self._endMousePos:
            return
        pen = QtGui.QPen()
        pen.setStyle(QtCore.Qt.PenStyle.DotLine) # 设置为虚线
        pen.setColor(QtCore.Qt.GlobalColor.green) # 设置线条颜色
        pen.setWidth(2) # 设置线条宽度
        p.setPen(pen)
        drawrect = QtCore.QRect(self._beginMousePos.x()
                         , self._beginMousePos.y()
                         , self._endMousePos.x() - self._beginMousePos.x() 
                         , self._endMousePos.y() - self._beginMousePos.y())
        p.drawRect(drawrect)
        selectrect = self.getSelectArea()
        if selectrect:
            # 保存的已经是原图选框的大小了，不要再转了
            nx, ny, nw, nh = convertDisplay2Original((selectrect.x(), selectrect.y(), selectrect.width(), selectrect.height())
                                                     , (self.width(), self.height())
                                                    , (self._image.width(), self._image.height()))
            nrect = QtCore.QRect(nx, ny, nw, nh)
            self._lastSelectedArea = nrect
            img = self._image.scaled(self.width(), self.height()
                                     , QtCore.Qt.AspectRatioMode.IgnoreAspectRatio
                                     , QtCore.Qt.TransformationMode.SmoothTransformation)
            selectimg = img.copy(selectrect)
            drawpos = QtCore.QPoint(selectrect.x(), selectrect.y())
            p.drawImage(drawpos, selectimg)
        return True

class CtrlWindowUI(QtWidgets.QMainWindow):

    def __init__(self, parent = None):
        super().__init__(parent)
        self._savescreenshot: bool = False
        self._recording: bool = False
        self._lastrecordingtime = 0
        self._projectname = ""

        self.setWindowTitle("AxisHomeCenter")
        self.bottomDockWidget = QtWidgets.QDockWidget(self)
        self.console = QtWidgets.QPlainTextEdit()
        common.globaldefines.g_DebugPrintHook = self.HookDebugPrint
        self.bottomDockWidget.setWidget(self.console)
        self.addDockWidget(QtCore.Qt.DockWidgetArea.BottomDockWidgetArea, self.bottomDockWidget)
        self.bottomDockWidget.resize(self.width(), 200)


        self.rightDockWidget = QtWidgets.QDockWidget(self)
        self.addDockWidget(QtCore.Qt.DockWidgetArea.RightDockWidgetArea, self.rightDockWidget)

        self.setLayout(QtWidgets.QHBoxLayout())
        self._displaywidget = DisplayWidget()
        self.setCentralWidget(self._displaywidget)
        #self.layout().addWidget(self._displaywidget)
        self._displaywidget.resize(896, 504)
        
        self._groupbox = QtWidgets.QGroupBox()
        #self._groupbox.setFixedWidth(200)
        #self.layout().addWidget(self._groupbox)
        self.rightDockWidget.setWidget(self._groupbox)
        self._groupbox.setLayout(QtWidgets.QVBoxLayout())

        self._scriptcomboBox = QtWidgets.QComboBox()
        self._groupbox.layout().addWidget(self._scriptcomboBox)

        self._executeflowBtn = QtWidgets.QPushButton("Execute")
        self._groupbox.layout().addWidget(self._executeflowBtn)

        self._makescreenshotBtn = QtWidgets.QPushButton("Screenshot")

        self._editMode = QtWidgets.QRadioButton("Edit")
        self._readonlyMode = QtWidgets.QRadioButton("Normal")
        self._readonlyMode.setChecked(True)

        self._groupbox.layout().addWidget(self._makescreenshotBtn)
        self._groupbox.layout().addWidget(self._editMode)
        self._groupbox.layout().addWidget(self._readonlyMode)

        #self._groupbox.layout().addItem(QtWidgets.QSpacerItem(20, 40, vData=QtWidgets.QSizePolicy.Policy.Expanding))
        
        self.labels = QtWidgets.QTreeWidget(self)
        self.labels.setHeaderLabels(["Label", "Area"])
        self._groupbox.layout().addWidget(self.labels)

        self.removeItemBtn = QtWidgets.QPushButton("Remove")
        self.recordDataBtn = QtWidgets.QPushButton("Record")

        self._groupbox.layout().addWidget(self.removeItemBtn)
        self._groupbox.layout().addWidget(self.recordDataBtn)

        self._title = ""

        self.initButton()

    def initButton(self):
        self._executeflowBtn.clicked.connect(self.onExecuteBtnClicked)
        self._makescreenshotBtn.clicked.connect(self.onMakeScreenshotBtnClicked)
        self._editMode.clicked.connect(self.onModeSwitch)
        self._readonlyMode.clicked.connect(self.onModeSwitch)
        self.removeItemBtn.clicked.connect(self.onRemoveLabel)
        self.recordDataBtn.clicked.connect(self.onRecordDataBtn)
        self._displaywidget.signalNewLabel.connect(self.onNewLabel)

    def HookDebugPrint(self, txt):
        self.console.appendPlainText(f"{txt}")

    def refreshLabelsUI(self):
        self.labels.clear()
        for info in labelmanager.g_Instance._labels.get(self._title, []):
            self.onNewLabel(info[0])

    def onNewLabel(self, name):
        info = labelmanager.g_Instance.getLabel(self._title, name)
        item = QtWidgets.QTreeWidgetItem()
        item.setText(0, info[0])
        item.setText(1, f"{info[1]},{info[2]} {info[3]}x{info[4]}")
        self.labels.addTopLevelItem(item)

    def onRemoveLabel(self):
        item = self.labels.currentItem()
        if not item:
            return
        labelname = item.text(0)
        labelmanager.g_Instance.removeLabel(self._title, labelname)
        idx = self.labels.indexOfTopLevelItem(item)
        if idx >= 0:
            self.labels.takeTopLevelItem(idx)

    def onRecordDataBtn(self):
        if not self._recording:
            self.recordDataBtn.setText("Stop Record")
            self._recording = True
        else:
            self.recordDataBtn.setText("Record")
            self._recording = False

    def onModeSwitch(self):
        if self._readonlyMode.isChecked():
            self._displaywidget.setReadOnlyMode(True)
        else:
            self._displaywidget.setReadOnlyMode(False)

    def onExecuteBtnClicked(self):
        DebugPrint("asdasd")
        flowid = self._scriptcomboBox.currentData()
        if not flowid in application.Application.instance().getThread("FlowManager")._flows:
            DebugPrint(f"{flowid} not found.")
            return
        DebugPrint(f"execute flow: {self._scriptcomboBox.currentText()}")
        application.g_App.pushMessage("FlowManager", "ExecuteFlow", flowid)

    def onMakeScreenshotBtnClicked(self):
        self._savescreenshot = True
        DebugPrint("???")

    def updateScreen(self, img):
        self._displaywidget.updateScreen(img)
        if self._savescreenshot:
            self._savescreenshot = False
            pixmap = QtGui.QPixmap(img)
            pixmap.save("screenshot.png", "png")
            DebugPrint("save screenshot.")
        if self._recording and time.time() * 1000 - self._lastrecordingtime > 1000:
            timestamp = GetTimeStamp()
            if not os.path.exists("recoding/"):
                os.makedirs("recoding")
            pixmap = QtGui.QPixmap(img)
            pixmap.save(f"recoding/{self._projectname}-{timestamp}.png", "png")

    